import argparse
from pathlib import Path
from auto_optimizer import OnnxGraph


def get_sub_structures(onnx_path):
    g = OnnxGraph.parse(onnx_path)
    matmul_nodes = g.get_nodes('MatMul')
    ini_names = set([node.name for node in g.get_nodes('Initializer')])

    id2name = {}
    for node in matmul_nodes:
        flag = False
        for input_name in node.inputs:
            if input_name in ini_names:
                flag = True
                break
        if not flag:
            continue
        id = int(node.name.replace('MatMul_', ''))
        id2name[id] = node.name

    names = [e[1] for e in sorted(id2name.items(), key=lambda x: x[0])]
    units = [names[s: s+6] for s in range(0, len(names)-1, 6)]
    assert len(units) == 24

    sub_list = [3, 3, 67, 15, 11]
    for i, unit_names in enumerate(units):
        last_id = int(unit_names[0].replace('MatMul_', ''))
        if i != 0:
            assert last_id-int(units[i-1][0].replace('MatMul_', '')) == 114
        for j, name in enumerate(unit_names[1:]):
            id = int(name.replace('MatMul_', ''))
            assert id-last_id == sub_list[j], f'{id}-{last_id} != {sub_list[j]}'
            last_id = id

    return units


def create_config(units, unit_indices, save_config=None, skip_tail=True):
    if not save_config:
        idx_str = '-'.join([str(idx) for idx in unit_indices])
        save_config = f'./quant_unit{unit_idx}.cfg'
    f = open(save_config, 'w', encoding='utf-8')
    for i, unit_names in enumerate(units):
        if i in set(unit_indices):
            print("quantize:", unit_names)
            continue
        for name in unit_names:
            f.write(f'skip_layers : "{name}"\n')
    if skip_tail:
        f.write(f'skip_layers : "MatMul_3171"\n')
    f.close()


def main():
    parser = argparse.ArgumentParser(
        description='create config of quantization.')
    parser.add_argument('onnx', type=str, help='input onnx.')
    parser.add_argument('config', type=str, help='output onnx.')
    args = parser.parse_args()

    sub_structures = get_sub_structures(args.onnx)
    indices = list(range(0, 10)) + list(range(15, 24))
    create_config(sub_structures, indices, args.config)


if __name__ == "__main__":
    main()
